/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#include "expression_impl.h"

#include <queue>
#include <symengine/assumptions.h>
#include <symengine/functions.h>
#include <symengine/simplify.h>
#include <symengine/integer.h>
#include <symengine/real_double.h>

#include "expr_print_manager.h"
#include "const_values.h"
#include "common/util/mem_utils.h"
namespace ge {
namespace expression {
ExpressionImpl::~ExpressionImpl() {}

std::string ExpressionImpl::Str(const StrType type) const {
  if (type == StrType::STR_CPP) {
    if (SymEngine::is_a<SymEngine::Rational>(*sym_expr_)) {
      const SymEngine::Rational &x = SymEngine::down_cast<const SymEngine::Rational &>(*sym_expr_);
      auto dens = x.get_den();
      auto nums = x.get_num();
      return "(double)(" + nums->__str__() +")/(double)(" + dens->__str__() + ")";
    }
    if ((GetExprType() == ExprType::EXPR_OPERATION) &&
        (GetOpType() != OperationType::OP_NONE)) {
      auto printer = ExprManager::GetInstance().GetPrinter(GetOpType());
      return printer(GetArgs());
    }
    return sym_expr_->__str__();
  }
  return "";
}

ExpressionImplPtr ExpressionImpl::Replace(const std::map<ExpressionImplPtr, ExpressionImplPtr> &replace_vars) const {
  SymEngine::map_basic_basic sym_replace_vars;
  for (const auto &sym_expr_impl_ptr_item : replace_vars) {
    sym_replace_vars.emplace(sym_expr_impl_ptr_item.first->sym_expr_,
                             sym_expr_impl_ptr_item.second->sym_expr_);
  }
  SymEngineExprPtr replaced_expr = sym_expr_->subs(sym_replace_vars);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(replaced_expr);
}

ExpressionImplPtr ExpressionImpl::Simplify() const {
  SymEngineExprPtr simplified_expr = SymEngine::simplify(sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(simplified_expr);
}

bool ExpressionImpl::ContainVar(const ExpressionImplPtr &e) const {
  if (e->sym_expr_->get_args().size() != 0u) {
    return false;
  }
  for (auto &arg : GetPrimarySymArgs()) {
    if (SymEngine::eq(*arg, *(e->sym_expr_))) {
      return true;
    }
  }
  return false;
}

bool ExpressionImpl::operator==(const ExpressionImpl &e) const {
  return SymEngine::eq(*sym_expr_, *e.sym_expr_);
}

std::vector<SymEngineExprPtr> ExpressionImpl::GetPrimarySymArgs() const {
  std::vector<SymEngineExprPtr> single_args;
  std::queue<SymEngineExprPtr> args_queue;
  args_queue.push(sym_expr_);
  while (!args_queue.empty()) {
    auto arg = args_queue.front();
    args_queue.pop();
    auto arg_elements = arg->get_args();
    if (arg_elements.empty()) {
      // 过滤非变量
      if (SymEngine::is_a_Number(*arg)) {
        continue;
      }
      // 去重
      bool has_find = false;
      for (auto &single_arg : single_args) {
        if (SymEngine::eq(*single_arg, *arg)) {
          has_find = true;
          break;
        }
      }
      if (!has_find) {
        single_args.emplace_back(arg);
      }
    } else {
      for (auto &element : arg_elements) {
        args_queue.push(element);
      }
    }
  }
  return single_args;
}

std::vector<ExpressionImplPtr> ExpressionImpl::GetPrimaryArgs() const {
  std::vector<ExpressionImplPtr> args_list;
  for (auto &arg : GetPrimarySymArgs()) {
    args_list.emplace_back(ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(arg));
  }
  return args_list;
}

std::vector<SymEngineExprPtr> ExpressionImpl::GetArgs() const {
  return sym_expr_->get_args();
}

double ExpressionImpl::GetIntegerResult(const SymEngine::Integer &integer_expr) {
  if (integer_expr.is_zero()) {
    return 0;
  } else if (integer_expr.is_positive()) {
    return static_cast<double>(integer_expr.as_uint());
  }
  return static_cast<double>(integer_expr.as_int());
}

bool ExpressionImpl::GetResult(double &result) {
  if (SymEngine::is_a<SymEngine::Integer>(*sym_expr_)) {
    const auto &integer_expr = SymEngine::down_cast<const SymEngine::Integer&>(*sym_expr_);
    result = GetIntegerResult(integer_expr);
    return true;
  }
  if (SymEngine::is_a<SymEngine::Rational>(*sym_expr_)) {
    const auto &rational_expr = SymEngine::down_cast<const SymEngine::Rational&>(*sym_expr_);
    result = GetIntegerResult(*(rational_expr.get_num())) /
      GetIntegerResult(*(rational_expr.get_den()));
    return true;
  }
  if (SymEngine::is_a<SymEngine::RealDouble>(*sym_expr_)) {
    const auto &real_double_expr = SymEngine::down_cast<const SymEngine::RealDouble&>(*sym_expr_);
    result = real_double_expr.as_double();
    return true;
  }
  return false;
}

OperationType ExpressionImpl::GetOpType() const {
  if (SymEngine::is_a<SymEngine::Add>(*sym_expr_)) {
    return OperationType::OP_ADD;
  }
  if (SymEngine::is_a<SymEngine::Mul>(*sym_expr_)) {
    return OperationType::OP_MUL;
  }
  if (SymEngine::is_a<SymEngine::Max>(*sym_expr_)) {
    return OperationType::OP_MAX;
  }
  if (SymEngine::is_a<SymEngine::Min>(*sym_expr_)) {
    return OperationType::OP_MIN;
  }
  if (SymEngine::is_a<SymEngine::Pow>(*sym_expr_)) {
    return OperationType::OP_POW;
  }
  if (SymEngine::is_a<SymEngine::Log>(*sym_expr_)) {
    return OperationType::OP_LOG;
  }
  if (SymEngine::is_a<SymEngine::Ceiling>(*sym_expr_)) {
    return OperationType::OP_CEIL;
  }
  return OperationType::OP_NONE;
}

ExprType ExpressionImpl::GetExprType() const {
  if (SymEngine::is_a_Number(*sym_expr_)) {
    return ExprType::EXPR_CONSTANT;
  }
  if (SymEngine::is_a<SymEngine::Symbol>(*sym_expr_)) {
    return ExprType::EXPR_VARIABLE;
  }
  return ExprType::EXPR_OPERATION;
}

ExpressionImpl::ExpressionImpl(const int32_t const_value) :
  sym_expr_(SymEngine::integer(const_value)) {}

ExpressionImpl::ExpressionImpl(const int64_t const_value) :
  sym_expr_(SymEngine::integer(const_value)) {}

ExpressionImpl::ExpressionImpl(const uint32_t const_value) :
  sym_expr_(SymEngine::integer(const_value)) {}

ExpressionImpl::ExpressionImpl(const uint64_t const_value) :
  sym_expr_(SymEngine::integer(const_value)) {}

ExpressionImpl::ExpressionImpl(const double const_value) :
  sym_expr_(SymEngine::real_double(const_value)) {}

ExpressionImpl::ExpressionImpl(const std::string &name) :
  sym_expr_(SymEngine::symbol(name)) {}

ExpressionImpl::ExpressionImpl(const SymEngineExprPtr &sym_expr) :
  sym_expr_(sym_expr) {}

ExpressionImplPtr Add(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::add(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Sub(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::sub(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Mul(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::mul(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Div(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::div(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Max(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::max({a->sym_expr_, b->sym_expr_});
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Min(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::min({a->sym_expr_, b->sym_expr_});
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Pow(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::pow(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Log(const ExpressionImplPtr &a) {
  SymEngineExprPtr sym_expr = SymEngine::log(a->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Log(const ExpressionImplPtr &a, const ExpressionImplPtr &b) {
  SymEngineExprPtr sym_expr = SymEngine::log(a->sym_expr_, b->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}

ExpressionImplPtr Ceiling(const ExpressionImplPtr &a) {
  SymEngineExprPtr sym_expr = SymEngine::ceiling(a->sym_expr_);
  return ExpressionImpl::CreateExpressionImpl<const SymEngineExprPtr&>(sym_expr);
}
std::ostream &operator<<(std::ostream &os, const ExpressionImpl &e) {
  os << e.Str();
  return os;
}
}  // namespace expression
}  // namespace ge
